#!/usr/bin/env python3

import json
from .tool_register import ToolRegistry
from .tool_functions import tool_registry

def get_json_schema_type(python_type):
    type_map = {
        'str': "string",
        'int': "number",
        'float': "number",
        'dict': "object",
        'list': "array",
        'bool': "boolean",
        'NoneType': "null"
    }
    return type_map.get(python_type, "string")

def generate_schema(tool_registry):
    tool_names = list(tool_registry.tools.keys())

    tool_input_schemas = {
        tool_name: {
            "type": "object",
            "properties": {
                arg_name: {
                    "type": get_json_schema_type(arg_info["type"])
                }
                for arg_name, arg_info in tool.args.items()
            },
            "required": list(tool.args.keys())
        }
        for tool_name, tool in tool_registry.tools.items()
    }

    schema = {
        "title": "Response",
        "type": "object",
        "properties": {
            "think_step_by_step": {
                "type": "string"
            },
            "response": {
                "type": "string"
            },
            "tool_call": {
                "type": "array",
                "items": {
                    "oneOf": [
                        {
                            "type": "object",
                            "properties": {
                                "tool": {"const": tool_name},
                                "tool_input": tool_input_schemas[tool_name]
                            },
                            "required": ["tool", "tool_input"]
                        } for tool_name in tool_names
                    ]
                },
                "minItems": 0,
                "maxItems": 1
            }
        },
        "required": ["response", "think_step_by_step", "tool_call"]
    }
    return json.dumps(schema, indent=4)


if __name__ == "__main__":
    schema = generate_schema(tool_registry)
    print(schema)
